# coding=utf-8
import time
from openai import OpenAI, AsyncOpenAI
from langchain_core.output_parsers import StrOutputParser
from utils.generate_token import api_key
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
import json
import tiktoken



# AI回复
class AI_Answer:

    def __init__(self):
        self.llm = ChatZhipuAI(
            api_key=api_key,
            model_name="glm-4-plus",
            streaming=True,
            temperature=0.01,
            top_p=0.01,
        )

        self.client = AsyncOpenAI(
            api_key="5db64dc0-73d7-4139-b755-5bd5483327b4",  # lz
            base_url="https://ark.cn-beijing.volces.com/api/v3",
        )

        # 使用本地大模型
        self.local_client = AsyncOpenAI(
            base_url="http://192.168.2.186:11434/v1",
            api_key="",
        )


    # 统计token消耗
    def count_tokens(self, text: str, model: str = 'cl100k_base'):
        """
        统计文本的token数量

        :param text: 输入的文本
        :param model: 模型名称
        :return: token数量
        """
        try:
            encoding = tiktoken.encoding_for_model(model)
        except KeyError:
            encoding = tiktoken.get_encoding('cl100k_base')
        return len(encoding.encode(text))


    # 智谱AI回答
    async def get_ai_answer(self, prompt: str):

        try:
            query_message = prompt
            chat_prompt = ChatPromptTemplate.from_messages(
                [("system", "你是一个证券交易所中审核上市文件的专家。"),
                 ("human", "{query_message}")]
            )
            chain = chat_prompt | self.llm | StrOutputParser()
            async for token in chain.astream({"query_message": query_message}):
                yield token
        except Exception as e:
            yield f"C5 error : {e}"

    # DeepSeek-R1 AI回答
    async def get_ai_answer_R1(self, prompt: str, output_json: bool = False):
        """
        向AI模型发送请求并获取流式回答

        :param prompt: 输入的提示文本
        :param output_json: 是否要求输出JSON格式
        :return: 异步生成器，产生AI回答的 chunks
        """
        from utils.logger import logger

        if not prompt or not prompt.strip():
            logger.error("Prompt cannot be empty")
            yield json.dumps({'type': 'error', 'content': "C5 error: Prompt cannot be empty"}, ensure_ascii=False)
            return

        try:
            s_time = time.time()
            logger.info("----- Starting streaming request -----")
            logger.debug(f"Prompt: {prompt}")

            # 准备请求参数
            request_params = {
                "model": "ep-20250210135803-8r8cd",  # lz 满血
                "messages": [{"role": "user", "content": prompt}],
                "stream": True,
                "timeout": 300
            }

            # 如果需要JSON输出，添加response_format参数
            if output_json:
                request_params["response_format"] = {"type": "json_object"}

            # 发送请求
            stream = await self.client.chat.completions.create(**request_params)

            reasoning_content = ""
            content = ""
            async for chunk in stream:
                if hasattr(chunk.choices[0].delta, 'reasoning_content'):
                    reasoning_content += chunk.choices[0].delta.reasoning_content
                    yield json.dumps({
                        'type': 'thinking', 
                        'content': chunk.choices[0].delta.reasoning_content
                    }, ensure_ascii=False)
                else:
                    if chunk.choices[0].delta.content:
                        content += chunk.choices[0].delta.content
                        yield json.dumps({
                            'type': 'text', 
                            'content': chunk.choices[0].delta.content
                        }, ensure_ascii=False)

            e_time = time.time()
            logger.info(f"Streaming request completed in {e_time - s_time:.2f} seconds")
            logger.debug(f"Total content length: {len(content)}")

        except asyncio.TimeoutError:
            logger.error("C5 error: Request timed out")
            yield json.dumps({'type': 'error', 'content': "C5 error: Request timed out"}, ensure_ascii=False)
        except Exception as e:
            logger.error(f"C5 error: {str(e)}", exc_info=True)
            yield json.dumps({'type': 'error', 'content': f"C5 error: {str(e)}"}, ensure_ascii=False)

    # 调用本地部署的deepseek-r1模型
    async def get_ai_answer_R1_local(self, prompt: str, output_json: bool = False):
        """
        向本地部署的deepseek-r1模型发送请求并获取流式回答

        :param prompt: 输入的提示文本
        :param output_json: 是否要求输出JSON格式
        :return: 异步生成器，产生AI回答的 chunks
        """
        from utils.logger import logger

        if not prompt or not prompt.strip():
            logger.error("Prompt cannot be empty")
            yield json.dumps({'type': 'error', 'content': "C5 error: Prompt cannot be empty"}, ensure_ascii=False)
            return

        try:
            s_time = time.time()
            logger.info("----- Starting streaming request -----")
            logger.debug(f"Prompt: {prompt}")

            # 准备请求参数
            request_params = {
                "model": "modelscope.cn/unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF",  # 本地模型名称
                "messages": [{"role": "user", "content": prompt}],
                "stream": True,
                "timeout": 300
            }

            # 如果需要JSON输出，添加response_format参数
            if output_json:
                request_params["response_format"] = {"type": "json_object"}
            # 发送请求
            stream = await self.local_client.chat.completions.create(**request_params)
            reasoning_content = ""
            content = ""
            async for chunk in stream:
                if hasattr(chunk.choices[0].delta, 'reasoning_content'):
                    reasoning_content += chunk.choices[0].delta.reasoning_content
                    yield json.dumps({
                        'type': 'thinking', 
                        'content': chunk.choices[0].delta.reasoning_content
                    }, ensure_ascii=False)
                else:
                    if chunk.choices[0].delta.content:
                        content += chunk.choices[0].delta.content
                        yield json.dumps({
                            'type': 'text', 
                            'content': chunk.choices[0].delta.content
                        }, ensure_ascii=False)

            e_time = time.time()
            logger.info(f"Streaming request completed in {e_time - s_time:.2f} seconds")
            logger.debug(f"Total content length: {len(content)}")

        except asyncio.TimeoutError:
            logger.error("C5 error: Request timed out")
            yield json.dumps({'type': 'error', 'content': "C5 error: Request timed out"}, ensure_ascii=False)
        except Exception as e:
            logger.error(f"C5 error: {str(e)}", exc_info=True)
            yield json.dumps({'type': 'error', 'content': f"C5 error: {str(e)}"}, ensure_ascii=False)


    
if __name__ == '__main__':
    import asyncio
    prompt = "# CONTEXT #\n<content_text>与</content_text>之中的为招股说明书中的某一段文本切片。\n<query_dict>与</query_dict>之中的为提问字典；Key为索引，Value为一个字典{\"依据材料\": \"材料内容\", \"请发行人说明/披露\": \"问询问题\"}。<query_dict>与</query_dict>之中的每一个Value是一个dict。每一个dict代表交易所的历史审核问询中的最相似的某一个问询。其中，”依据材料“key的value是历史最相似问询所依据的招股说明书中的材料；”问询问题“key的value是其针对”依据材料“对发行人所提出的审核问询。\n招股说明书：\n<content_text>\n发行人与齐齐哈尔博实畜牧服务有限公司就支持和牛养殖事宜签署了《支持和牛养殖协议》，约定发行人为齐齐哈尔博实畜牧服务有限公司提供6,000万元额度的专项资金支持，用于和牛养殖服务；齐齐哈尔博实畜牧服务有限公司以其拥有的牛只对上述专项资金支持额度提供抵押担保，齐齐哈尔博实畜牧服务有限公司的直接股东彭辉及关联方龙江县彭辉肉牛养殖场、龙江黑牛牧业有限公司梅里斯分公司提供连带责任保证。\n</content_text>;\n提问字典：\n<query_dict>\n{}\n</query_dict>。\n# OBJECTIVE #\n从交易所的角度，基于<query_dict>与</query_dict>之中的历史相似问询问题，生成相应的<content_text>与</content_text>之间切片的最有可能被问询的问题。\n# STYLE #\n请以交易所IPO业务的外审人员的风格来进行问询，只输出”请发行人说明/披露“的内容。\n# TONE #\n正式。\n# AUDIENCE #\n我们的受众是投行IPO业务的内审人员和交易所IPO业务的外审人员。\n# RESPONSE #\n请严格遵循以上要求和下面输出示例的格式输出，不要输出“请发行人说明/披露”相关的标题!!!\n请严格遵循以上要求和下面输出示例的格式输出，不要输出“请发行人说明/披露”相关的标题!!!\n\n输出示例：\n(1)...；\n(2)...；\n......\n(n)...。\n"
    async def main():
        AI = AI_Answer()
        async for response in AI.get_ai_answer_R1_local(prompt):
            print(response)
    asyncio.run(main())
